#!/usr/bin/env python
# -*- coding: UTF-8 -*-
import redis

# 连接池配置
pool = redis.ConnectionPool(
    host='localhost',
    port=6379,
    max_connections=10,
)

r = redis.Redis(
    host='localhost',   # Redis服务器IP
    port=6379,          # 端口
    db=0,               # 数据库索引（0-15）
    password='123456', # 密码（若无密码可省略）
    decode_responses=True,  # 自动解码返回值为字符串（默认bytes）
    connection_pool=pool # 连接池
)

def testPipe():
    # 普通命令 (多次网络请求，每次执行都是一次请求)
    r.set("counter1", 0)
    r.incr("counter1") # 自增
    r.incr("counter1")
    print(f'counter1: {r.get("counter1")}')
    # 管道方式 (单词网络请求)
    pipe = r.pipeline()
    pipe.set("counter2", 0)
    pipe.incr("counter2")
    pipe.incr("counter2")
    pipe.execute()
    print(f'counter2: {r.get("counter2")}')

def testMulti(): # 事务
    # 监视键，防止被其他客户端修改
    r.set("balance", 100)
    r.watch("balance")

    pipe = r.pipeline()
    try:
        pipe.multi() # 开启事务
        pipe.decrby("balance", 20)
        pipe.incrby("savings", 20)
        pipe.execute() # 提交事务
        print(f'balance: {r.get("balance")}')
        print(f'savings: {r.get("savings")}')
    except redis.WatchError:
        print("事务执行失败（balance被修改）")

def testPubSub():
    # 实现消息的实时广播
    # 发布端:
    r.publish("news_channel", "Hello, Redis Pub/Sub!")
    # 接收端:
    pubsub = r.pubsub()
    pubsub.subscribe("news_channel")

    for message in pubsub.listen():
        if message["type"] == "message":
            print(f"收到消息: {message['data'].decode()}")

def testLua():
    # 编写脚本
    luaScript = """
        local key = KEYS[1]  -- 获取传入的键（用户的请求计数）
        local limit = tonumber(ARGV[1])  -- 获取请求限制的最大次数，转换为数字类型
        local curr = tonumber(redis.call('GET', key) or "0")  -- 获取当前的请求次数，若键不存在则默认为 0
        if curr + 1 > limit then  -- 检查当前请求次数加 1 是否超过限制
            return 0  -- 如果超过限制，返回 0，表示拒绝请求
        else 
            redis.call('INCR', key)  -- 否则，增加该键的值，表示成功请求
        end
    """
    # 调用脚本
    result = r.eval(luaScript, 1, "rate_limit:user10", 100)
    print("允许访问" if result else "超过限制")

def main():
    # testPipe()
    # testMulti()
    # testPubSub()
    testLua()


if __name__ == "__main__":
    main()